/*
 * Copyright (c) China Telecom Cloud Technology Co., Ltd. 2024-2025. All rights reserved.
 * Created by l00841256 on 25-3-31.
 */
#include "KVCacheService.h"
#include "Common.h"
#include "MemoryPool.h"
#include <map>
#include <string>
#include <vector>
#ifdef USE_DYNAMO
#include "DynamoPublisher.h"
#endif

void KVCacheServiceImpl::buildPrefetchPrefixCacheTaskFuture(
                            std::vector<std::future<int32_t>> &taskFutures, 
                            MemoryPool *pool,
                            HashTask getCacheTask, 
                            int32_t tpNum, 
                            CacheInfo cacheInfo) {
    // range: [startLayer, endLayer)
    uint64_t startOffset = cacheInfo.layerCacheSize * cacheInfo.startLayer;
    uint64_t length      = cacheInfo.layerCacheSize * (cacheInfo.endLayer - \
                                                 cacheInfo.startLayer);

    taskFutures.emplace_back(getThreadPool.submit([this, pool, getCacheTask,
                                                   startOffset, length, tpNum,
                                                   cacheInfo]() -> int32_t 
        {    // lambda func
            //! 如果在statusMap中找到了getCacheTask的prefixCacheKey，
            // 就不需要再支持从存储捞取相关哈希对应块的KV Cache到内存池
            if (!statusManager.shouldExecute(getCacheTask)) {
                LOG_WARN_FMT("should not execute! segmentId: {}",
                            getCacheTask.blockHashId.hashId);
                return KVC_OK;
            }

            // 尝试分配1个segment
            std::vector<BlockHashId> blockHashId(1, getCacheTask.blockHashId);
            // 
            std::vector<Segment *> segments;
            // TODO(): 这里为什么writeInRecordSize直接赋值为1
            int32_t ret = pool->allocateSegments(blockHashId, ALLOC_POLICY_LRU,
                                                segments, 1, cacheInfo);

            //! Segment: TP , START LAYER - END LAYER , BLCOK-size
            //!  File: TP ,  LAYERs , BLCOK-size
            if (ret != KVC_OK) {
                statusManager.recordFailure(getCacheTask);
                LOG_WARN_FMT("allocate segment failed! segmentId: {}",
                            getCacheTask.blockHashId.hashId);
                return KVC_ERR_NOMEM;
            }
            Segment *segment = segments[0];

            // 尝试下盘读
            std::string objId = std::to_string(getCacheTask.blockHashId.hashId) +
                                FILE_NAME_SEPARATOR + std::to_string(tpNum);
            
            //! TODO(yangxianpku: 2025.08.21): 参数为什么是这样？
            ret = storageClient.get(objId, *segment, startOffset, startOffset,
                                    length);
            if (ret != KVC_OK) {
                statusManager.recordFailure(getCacheTask);
                LOG_WARN_FMT("read fail! segmentId: {}", segment->getSegmentId());
                pool->releaseSegment(segment->getSegmentId());
                return KVC_ERR_INTERNAL;
            }

            // 下盘读成功，标记segment为只读
            pool->sealSegment(getCacheTask.blockHashId.hashId);
            segment->setIsSaved(true);
            return KVC_OK;
        }
    ));

    return;
}




void KVCacheServiceImpl::processPrefetchPrefixCacheRequest(
            const kvcache::PrefetchPrefixCacheRequest *request, 
            MemoryPool *pool,
            std::vector<HashTask> &getCacheTasks, 
            std::vector<int64_t> &allHashIds) {
    for (const auto &entry : request->entries()) {   // 遍历每个句子
        // 一个句子希望检索的所有BlockHashId
        std::vector<BlockHashId> blockHashIds;

        for (auto blkHid : entry.blockhashids()) {
            char tmpVerifyHashId[VERIFY_HASH_LENGTH];

            // TODO(yangxianpku: 2025.08.21): 更换为copy形式更高效
            for (int i = 0; i < VERIFY_HASH_LENGTH; i++) {
                tmpVerifyHashId[i] = blkHid.verifyhashid()[i];
            }
            blockHashIds.emplace_back(blkHid.hashid(), blkHid.parenthashid(),
                                      blkHid.localhashid(), tmpVerifyHashId);
        }

        Segment *segment;

        // 本句子未命中的哈希块
        std::vector<BlockHashId> unhitBlockHashIds;  // 未命中的BlockHashId

        // TODO(): 移除代码冗余
        std::vector<int64_t>     unhitHashIds;       // 未命中的哈希
        for (int32_t index = 0; index < blockHashIds.size(); ++index) {
            allHashIds.push_back(blockHashIds[index].hashId);

            // 查看特定哈希是否存在内存池中
            int32_t ret = pool->getSegmentByID(blockHashIds[index].hashId, 
                                            &segment
                                        );
            if (ret != KVC_OK) {
                unhitBlockHashIds.push_back(blockHashIds[index]);
                unhitHashIds.push_back(blockHashIds[index].hashId);
                continue;
            }

            segment->setAccessTime(getCurrentTimestampUs());
            // 该序列的全部哈希块都能命中
            if (index == blockHashIds.size() - 1) {
                LOG_INFO_FMT("fully hit dram, block num:{}",
                             blockHashIds.size());
            }
        }

        //! 如果存在未命中的，将未命中的哈希值一起进行一次哈希
        if (!unhitBlockHashIds.empty()) {
            size_t prefixCacheKey = VectorHash<int64_t>{}(unhitHashIds);

            for (uint32_t i = 0; i < unhitBlockHashIds.size(); ++i) {
                HashTask hashTask(prefixCacheKey, unhitBlockHashIds[i],
                                  nullptr);
                getCacheTasks.emplace_back(hashTask);
            }
        }
    }

    return;
}



grpc::Status KVCacheServiceImpl::PrefetchPrefixCache(
        grpc::ServerContext *context,
        const kvcache::PrefetchPrefixCacheRequest *request,
        kvcache::PrefetchPrefixCacheResponse *response) {
    int32_t ret          = KVC_OK;
    uint64_t start_total = getCurrentTimestampUs(); // 总入口时间

    do {
        // 构造CacheInfo耗时
        uint64_t stage_start = getCurrentTimestampUs();
        CacheInfo cacheInfo{request->cacheinfo().startlayer(),
                            request->cacheinfo().endlayer(),
                            request->cacheinfo().layernum(),
                            request->cacheinfo().layercachesize()};
        LOG_DEBUG_FMT("CacheInfo init cost time : {} us",
                      getCurrentTimestampUs() - stage_start);
        int32_t tpNum = request->tpnum();

        MemoryPool *pool = nullptr;
        uint64_t begin_get_mempool = getCurrentTimestampUs();

        ret = getMemoryPool(request->memsize(), &pool);
        if (ret != KVC_OK) {
            LOG_ERROR_FMT("getMemoryPool failed, ret: {}", ret);
            break;
        }
        uint64_t after_get_mempool = getCurrentTimestampUs();
        LOG_DEBUG_FMT("Get Mempool cost time : {} us",
                      getCurrentTimestampUs() - stage_start);

        // 本次请求所有序列中未在Mempool中命中，需要尝试从存储捞取KV Cache的任务
        std::vector<HashTask> getCacheTasks; 

        // 本请求的所有句子所有块的哈希列表
        std::vector<int64_t> allHashIds;
        processPrefetchPrefixCacheRequest(request, 
                                        pool, 
                                        getCacheTasks,
                                        allHashIds
                                    );


        std::vector<std::future<int32_t>> taskFutures;
        taskFutures.reserve(getCacheTasks.size());

        uint64_t begin_prefetch_get_kvcache = getCurrentTimestampUs();

        // TODO(yangxianpku: 2025.08.21): 考虑预取的时候将本次已查询到在内存池中的块pin住
        for (auto getCacheTask : getCacheTasks) {
            buildPrefetchPrefixCacheTaskFuture(taskFutures, 
                                            pool, 
                                            getCacheTask,
                                            tpNum, 
                                            cacheInfo
                                        );
        }

        // 阻塞等待每个预取任务完成
        for (auto &taskFuture : taskFutures) {
            taskFuture.get();
        }
        uint64_t end_prefetch_get_kvcache = getCurrentTimestampUs();
        LOG_DEBUG_FMT("Prefetch Get KVCache cost time: {} us",
                      end_prefetch_get_kvcache - begin_prefetch_get_kvcache);

        // 将本次的KV 放在队列末尾，表示最近访问过
        pool->reorder(allHashIds);

        uint64_t end_reorder = getCurrentTimestampUs();
        LOG_DEBUG_FMT("Pool reorder cost time: {} us",
                      end_reorder - end_prefetch_get_kvcache);
    } while (0);

    uint64_t end_total = getCurrentTimestampUs(); // 总出口时间
    LOG_DEBUG_FMT("PrefetchPrefixCache costtime: {} us",
                  end_total - start_total);

    response->set_result(KVC_OK);
    return grpc::Status::OK;
}




void KVCacheServiceImpl::submitWriteSegmentTask(MemoryPool *pool,
                                                BlockHashId blockHashId,
                                                int32_t tpNum) {
    putThreadPool.submit([this, pool, blockHashId, tpNum]() {
        std::string objId = std::to_string(blockHashId.hashId) +
                            FILE_NAME_SEPARATOR + std::to_string(tpNum);
        if (storageClient.isExist(objId)) {
            LOG_INFO_FMT("File exist! hashId: {}", blockHashId.hashId);
            return;
        }

        Segment *segment;
        uint64_t begin_put_time = getCurrentTimestampUs();
        if (pool->getSegmentByID(blockHashId.hashId, &segment) == KVC_OK) {
            uint64_t startOffset =
                segment->getStartLayer() * segment->getLayerCacheSize();
            uint64_t length =
                (segment->getEndLayer() - segment->getStartLayer()) *
                segment->getLayerCacheSize();

            storageClient.put(objId, *segment, startOffset, startOffset,
                              length);
            segment->setIsSaved(true);
        }
        uint64_t end_put_time = getCurrentTimestampUs();
        LOG_DEBUG_FMT("put kvcache cost time: {} us",
                      end_put_time - begin_put_time);
    });

    return;
}



grpc::Status KVCacheServiceImpl::sealMemory(grpc::ServerContext *context,
                               const kvcache::sealMemoryRequest *request,
                               kvcache::sealMemoryResponse *response) {
    if (request->hashids().size() <= 0) {
        response->set_result(KVC_ERR_INVALID_ARG);
        LOG_INFO_FMT("sealMemory end: no hashId");
        return grpc::Status::OK;
    }

    int ret = KVC_OK;
    uint64_t start_total = getCurrentTimestampUs();
    do {
        std::vector<int64_t> blockHashIds;
        for (int32_t i = 0; i < request->hashids().size(); i++) {
            blockHashIds.push_back(request->hashids()[i]);
        }

        int32_t tpNum = request->tpnum();
        MemoryPool *pool = nullptr;

        ret = getMemoryPoolByHashId(blockHashIds[0], &pool);
        if (ret != KVC_OK) {
            break;
        }

        for (auto blockHashId : blockHashIds) {
            ret = pool->sealSegment(blockHashId);
            if (ret == KVC_OK && request->dosave()) {
                LOG_DEBUG("request dosave is True");
                submitWriteSegmentTask(pool, blockHashId, tpNum);
            }
        }
    } while (0);

    uint64_t end_total = getCurrentTimestampUs();
    LOG_DEBUG_FMT("sealMemory cost time: {} us", end_total - start_total);

    response->set_result(ret); // 示例返回值
    return grpc::Status::OK;
}



int32_t KVCacheServiceImpl::getMemoryPool(uint32_t memSize, MemoryPool **pool) {
    auto start = getCurrentTimestampUs();

    // TODO(yangxianpku： 2025.08.20)：这里为什么查找2次？？？
    //! 代码冗余
    auto it = _pools.find(memSize);
    if (it != _pools.end()) {
        *pool = it->second.get();
        LOG_INFO_FMT("Find already exists mempool.");
        return KVC_OK;
    }

    std::lock_guard<std::mutex> lock(_poolsMutex);

    it = _pools.find(memSize);
    if (it != _pools.end()) {
        LOG_INFO_FMT("Find already exists mempool.");
        *pool = it->second.get();
        return KVC_OK;
    }

    auto new_pool = std::make_unique<MemoryPool>();

    LOG_DEBUG_FMT("getMemoryPool call createMemoryPool");
    // TODO(yangxianpku 2025.08.26)： 这里传递一级指针不会有问题么
    const int32_t ret = createMemoryPool(MEMPOOL_SIZE, 
                                        memSize, 
                                        true, 
                                        false, 
                                        new_pool.get()
                                    );
    if (ret != KVC_OK) {
        LOG_ERROR_FMT("createMemoryPool failed! ret = {}", ret);
        return ret; // 自动释放 new_pool
    }

    //! C++17解构语法, 类似Python
    // emplace方法返回一个pair<iter, status>, iter表示插入的元素, status为插入状态
    auto [inserted_it, success] = _pools.emplace(memSize, 
                                        std::move(new_pool)
                                    );

    *pool = inserted_it->second.get();

    auto end = getCurrentTimestampUs();

    LOG_DEBUG_FMT("getMemoryPool cost time: {} us", end - start);
    return KVC_OK;
}




grpc::Status KVCacheServiceImpl::QueryPrefixCache(
            grpc::ServerContext *context,
            const kvcache::QueryPrefixCacheRequest *request,
            kvcache::QueryPrefixCacheResponse *response) {
    auto start   = getCurrentTimestampUs();
    auto hashIds = std::vector(request->hashids().begin(), 
                            request->hashids().end()
                        );
    if (hashIds.empty()) {
        LOG_INFO_FMT("hashIds empty!");
        return grpc::Status::OK;
    }

    // 获取memPool
    MemoryPool *pool = nullptr;
    int32_t ret      = getMemoryPoolByHashId(hashIds[0], &pool);
    if (ret != KVC_OK || !pool) {
        LOG_ERROR_FMT("getMemoryPoolByHashId failed! ret = {}", ret);
        return grpc::Status(grpc::StatusCode::NOT_FOUND,
                            "Cannot find memory pool");
    }

    response->set_memsize(pool->getMemSegSize());
    response->set_shmpath(pool->getShmFilePath());

    // 筛选segmentId，成功的返回其offset
    std::vector<Segment> segments = pool->readInOrder(hashIds);
    pool->updateCacheStatis(segments.size(), hashIds.size());
    if (segments.empty()) {
        LOG_INFO_FMT("read empty segments");
        return grpc::Status::OK;
    }

    LOG_INFO_FMT("Current request hit ratio : {:.2f}",
                 calculatePercentage(segments.size(), hashIds.size()));
    pool->printCacheStatis();

    // 包装结果
    Segment segment;
    for (uint32_t i = 0; i < segments.size(); ++i) {
        Segment &segment = segments[i];
        kvcache::SegmentAttr *attr = response->add_segmentattr();
        attr->set_segmentid(hashIds[i]);
        attr->set_offset(segment.getOffset());
        attr->set_writeinrecordsize(segment.getWriteInRecord());
        attr->set_startlayer(segment.getStartLayer());
        attr->set_endlayer(segment.getEndLayer());
        attr->set_layernum(segment.getLayerNum());
        attr->set_layercachesize(segment.getLayerCacheSize());
    }

    auto end = getCurrentTimestampUs();
    LOG_DEBUG_FMT("cost time : {} us", end - start);
    return grpc::Status::OK;
}




grpc::Status KVCacheServiceImpl::BatchQueryPrefixCache(
    grpc::ServerContext *context,
    const kvcache::BatchQueryPrefixCacheRequest *request,
    kvcache::BatchQueryPrefixCacheResponse *response) {
    auto start = getCurrentTimestampUs();
    auto queries =
        std::vector(request->queries().begin(), request->queries().end());

    MemoryPool *pool = nullptr;
    for (kvcache::QueryHashIds queryHashIds : queries) {
        auto hashIds = std::vector(queryHashIds.hashids().begin(),
                                   queryHashIds.hashids().end());

        if (hashIds.size() == 0) {
            kvcache::SegmentAttrList *attrList = response->add_attrlist();
            continue;
        }

        if (pool == nullptr) {
            int32_t ret = getMemoryPoolByHashId(hashIds[0], &pool);
            if (ret != KVC_OK || !pool) {
                LOG_ERROR_FMT("getMemoryPoolByHashId failed! ret : {}", ret);
                return grpc::Status(grpc::StatusCode::NOT_FOUND,
                                    "Cannot find memory pool");
            }
            response->set_memsize(pool->getMemSegSize());
            response->set_shmpath(pool->getShmFilePath());
        }

        // 筛选segmentId，成功的返回其offset
        std::vector<Segment> segments = pool->readInOrder(hashIds);
        pool->updateCacheStatis(segments.size(), hashIds.size());

        LOG_INFO_FMT("Current request hit ratio : {}%",
                     calculatePercentage(segments.size(), hashIds.size()));
        pool->printCacheStatis();

        // 包装结果
        Segment segment;
        kvcache::SegmentAttrList *attrList = response->add_attrlist();
        for (uint32_t i = 0; i < segments.size(); ++i) {
            Segment &segment = segments[i];
            kvcache::SegmentAttr *attr = attrList->add_attr();
            attr->set_segmentid(hashIds[i]);
            attr->set_offset(segment.getOffset());
            attr->set_writeinrecordsize(segment.getWriteInRecord());
            attr->set_startlayer(segment.getStartLayer());
            attr->set_endlayer(segment.getEndLayer());
            attr->set_layernum(segment.getLayerNum());
            attr->set_layercachesize(segment.getLayerCacheSize());
        }
    }
    auto end = getCurrentTimestampUs();
    LOG_DEBUG_FMT("cost time : {} us", end - start);
    return grpc::Status::OK;
}




grpc::Status KVCacheServiceImpl::AllocateMemory(
            grpc::ServerContext *context, 
            const kvcache::AllocateMemoryRequest *request,
            kvcache::AllocateMemoryResponse *response) {
    auto start = getCurrentTimestampUs(); // 记录函数开始时间

    // 反序列化
    std::vector<BlockHashId> blockHashIds;
    for (auto blockHashIdAttr : request->blockhashidattrs()) {
        char tmpVerifyHashId[VERIFY_HASH_LENGTH];
        for (int i = 0; i < VERIFY_HASH_LENGTH; i++) {
            tmpVerifyHashId[i] = blockHashIdAttr.verifyhashid()[i];
        }
        blockHashIds.emplace_back(
            blockHashIdAttr.hashid(), blockHashIdAttr.parenthashid(),
            blockHashIdAttr.localhashid(), tmpVerifyHashId);
    }

    uint32_t memSize = request->memsize();
    int32_t writeInRecordSize = request->writeinrecordsize();
    CacheInfo cacheInfo{
        request->cacheinfo().startlayer(), request->cacheinfo().endlayer(),
        request->cacheinfo().layernum(), request->cacheinfo().layercachesize()};

    MemoryPool *pool = nullptr;
    int32_t ret = getMemoryPool(memSize, &pool);
    if (ret != KVC_OK || !pool) {
        LOG_ERROR_FMT("getMemoryPool failed! ret = {}", ret);
        return grpc::Status(grpc::StatusCode::NOT_FOUND,
                            "Cannot find memory pool");
    }
    std::string shmPath = pool->getShmFilePath();
    response->set_shmpath(shmPath);

    std::vector<Segment *> segments;
    ret = pool->allocateSegments(blockHashIds, AllocPolicy::ALLOC_POLICY_LRU,
                                 segments, writeInRecordSize, cacheInfo);
    if (ret != 0) {
        // 如分配失败，则本次计算得到的block不下刷
        LOG_ERROR_FMT("allocateSegments failed! ret = {}", ret);
        return grpc::Status::OK;
    }

    for (int i = 0; i < segments.size(); i++) {
        kvcache::SegmentInfo *info = response->add_segments();
        info->set_poolid(0);
        info->set_segmentid(blockHashIds[i].hashId);
        info->set_offset(segments[i]->getOffset());
        info->set_memaddr(0);
        info->set_blocksize(memSize);
        info->set_allocatetime(0);
        info->set_isshare(true);
        info->set_isfilled(segments[i]->getIsFilled());
    }

    auto end = getCurrentTimestampUs(); // 记录函数结束时间
    LOG_DEBUG_FMT("cost time : {} us", end - start);
    return grpc::Status::OK;
}




grpc::Status KVCacheServiceImpl::QueryKVDistribute(
        grpc::ServerContext *context,
        const KVCtrl::QueryKVDistributeRequest *request,
        KVCtrl::QueryKVDistributeReply *response) {
    auto start = getCurrentTimestampUs();
    LOG_INFO_FMT("start");

    uint32_t memSize = request->mem_size();
    MemoryPool *pool = nullptr;
    int32_t ret      = getMemoryPool(memSize, &pool);
    if (ret != KVC_OK || !pool) {
        LOG_ERROR_FMT("getMemoryPool failed! ret = {}", ret);
        return grpc::Status(grpc::StatusCode::NOT_FOUND,
                            "Cannot find memory pool");
    }

    auto hashIds = pool->getAllCachedSegmentIds();
    for (auto hashId : hashIds) {
        response->add_session_ids(hashId.hashId);
    }

    auto end = getCurrentTimestampUs();
    LOG_DEBUG_FMT("cost time : {} us", end - start);
    return grpc::Status::OK;
}




void KVCacheServiceImpl::buildPrefetchTaskFuture(
    std::vector<std::future<int32_t>> &taskFutures, MemoryPool *pool,
    HashTask getCacheTask, int32_t tpNum) {
    taskFutures.emplace_back(getThreadPool.submit([this, pool, getCacheTask,
                                                   tpNum]() -> int32_t {
        if (!statusManager.shouldExecute(getCacheTask)) {
            LOG_INFO_FMT("should not execute! segmentId: {}",
                         getCacheTask.segment->getSegmentId());
            pool->releaseSegment(getCacheTask.segment->getSegmentId());
            return KVC_OK;
        }
        std::string objId = std::to_string(getCacheTask.blockHashId.hashId) +
                            FILE_NAME_SEPARATOR + std::to_string(tpNum);
        int32_t success =
            storageClient.get(objId, *(getCacheTask.segment), 0, 0,
                              pool->getMemSegSize()); // 假设的存储访问函数
        if (success != KVC_OK) {
            statusManager.recordFailure(getCacheTask);
            LOG_INFO_FMT("read fail! segmentId: {}",
                         getCacheTask.segment->getSegmentId());
            pool->releaseSegment(getCacheTask.segment->getSegmentId());
            return KVC_ERR_INTERNAL;
        }
        LOG_INFO_FMT("read success! hashId:{}, segmentId:{}",
                     getCacheTask.blockHashId.hashId,
                     getCacheTask.segment->getSegmentId());
        pool->sealSegment(getCacheTask.blockHashId.hashId);
        return KVC_OK;
    }));
    return;
}




int32_t KVCacheServiceImpl::processPrefetchRequest(
    const KVCtrl::PrefetchRequest *request, MemoryPool *pool,
    std::vector<HashTask> &getCacheTasks) {
    CacheInfo cacheInfo{
        request->cacheinfo().startlayer(), request->cacheinfo().endlayer(),
        request->cacheinfo().layernum(), request->cacheinfo().layercachesize()};
    const auto &hashIds = request->perfetch_ids();
    int32_t firstNotInMemoryIdx = hashIds.size();
    Segment *segment;
    for (int32_t index = 0; index < hashIds.size(); ++index) {
        int32_t ret = pool->getSegmentByID(hashIds[index], &segment);
        if (ret != KVC_OK && firstNotInMemoryIdx > index) {
            LOG_ERROR_FMT("getSegmentByID failed! ret = {}", ret);
            firstNotInMemoryIdx = index;
            break;
        }
    }

    std::vector<int64_t> unhitHashIds;
    if (firstNotInMemoryIdx < hashIds.size()) {
        unhitHashIds.insert(unhitHashIds.end(),
                            hashIds.begin() + firstNotInMemoryIdx,
                            hashIds.end());
    }

    if (!unhitHashIds.empty()) {
        std::vector<Segment *> allocatedSegments;
        std::vector<BlockHashId> unhitBlockHashIds;
        for (uint64_t unhitHashId : unhitHashIds) {
            unhitBlockHashIds.emplace_back(unhitHashId);
        }
        int32_t ret =
            pool->allocateSegments(unhitBlockHashIds, ALLOC_POLICY_LRU,
                                   allocatedSegments, 1, cacheInfo);
        if (ret != KVC_OK) {
            LOG_ERROR_FMT("allocateSegments failed, ret: {}", ret);
            return ret;
        }
        size_t prefixCacheKey = VectorHash<int64_t>{}(unhitHashIds);
        for (uint32_t i = 0; i < allocatedSegments.size(); ++i) {
            HashTask hashTask(prefixCacheKey, unhitBlockHashIds[i],
                              allocatedSegments[i]);
            getCacheTasks.emplace_back(hashTask);
        }
        perfixCacheSet.insert(prefixCacheKey);
    }
    return KVC_OK;
}




grpc::Status
KVCacheServiceImpl::Prefetch(grpc::ServerContext *context,
                             const KVCtrl::PrefetchRequest *request,
                             KVCtrl::PrefetchReply *response) {
    auto start = getCurrentTimestampUs();
    LOG_INFO_FMT("start");
    uint32_t memSize = request->mem_size();
    int32_t tpNum = request->tpnum();

    int32_t ret = KVC_OK;
    do {
        MemoryPool *pool = nullptr;
        ret = getMemoryPool(memSize, &pool);
        std::vector<HashTask> getCacheTasks;
        int32_t ret = processPrefetchRequest(request, pool, getCacheTasks);
        if (ret != KVC_OK) {
            break;
        }

        std::vector<std::future<int32_t>> taskFutures;
        taskFutures.reserve(getCacheTasks.size());
        uint32_t index = 0;
        for (auto getCacheTask : getCacheTasks) {
            buildPrefetchTaskFuture(taskFutures, pool, getCacheTask, tpNum);
        }
    } while (0);
    response->set_success(ret);

    auto end = getCurrentTimestampUs();
    LOG_DEBUG_FMT("cost time : {} us", end - start);
    return grpc::Status::OK;
}




grpc::Status
KVCacheServiceImpl::GetShmPath(grpc::ServerContext *context,
                               const kvcache::GetShmPathRequest *request,
                               kvcache::GetShmPathResponse *response) {
    auto start = getCurrentTimestampUs();
    MemoryPool *pool = nullptr;
    int32_t ret = getMemoryPool(request->memsize(), &pool);
    if (ret != KVC_OK || !pool) {
        LOG_ERROR_FMT("getMemoryPool failed");
        return grpc::Status(grpc::StatusCode::NOT_FOUND,
                            "Cannot find memory pool");
    }
    std::string shmPath = pool->getShmFilePath();
    response->set_shmpath(shmPath);

    auto end = getCurrentTimestampUs();
    LOG_DEBUG_FMT("cost time : {} us", end - start);
    return grpc::Status::OK;
}



grpc::Status KVCacheServiceImpl::CheckIfServiceReady(
    grpc::ServerContext *context,
    const kvcache::CheckIfServiceReadyRequest *request,
    kvcache::CheckIfServiceReadyResponse *response) {
    auto start = getCurrentTimestampUs();
    response->set_isready(isReady);

    auto end = getCurrentTimestampUs();
    LOG_DEBUG_FMT("cost time : {} us", end - start);
    return grpc::Status::OK;
}




grpc::Status KVCacheServiceImpl::CleanMemoryPool(
                    grpc::ServerContext *context,
                    const kvcache::CleanMemoryPoolRequest *request,
                    kvcache::CleanMemoryPoolResponse *response) {
    auto start = getCurrentTimestampUs();

    MemoryPool *pool = nullptr;
    int32_t ret = getMemoryPool(request->memsize(), &pool);
    if (ret != KVC_OK || !pool) {
        LOG_ERROR_FMT("getMemoryPool failed");
        return grpc::Status(grpc::StatusCode::NOT_FOUND,
                            "Cannot find memory pool");
    }

    pool->cleanMemPool();
    response->set_result(KVC_OK);


    // 清除statusMap中的所有项
    time_t zeroEntryTtl = 0;
    statusManager.cleanupExpiredEntries(zeroEntryTtl);

    auto end = getCurrentTimestampUs();
    LOG_DEBUG_FMT("cost time : {} us", end - start);
    return grpc::Status::OK;
}




int32_t KVCacheServiceImpl::getMemoryPoolByHashId(int64_t hashId,
                                                  MemoryPool **pPoolPtr) {
    auto start = getCurrentTimestampUs();
    if (_pools.size() == 1) {
        *pPoolPtr = _pools.begin()->second.get();
        return KVC_OK;
    }

    for (auto it = _pools.begin(); it != _pools.end(); it++) {
        MemoryPool *poolPtr = it->second.get();  // 拿到原始裸指针
        if (poolPtr->isInMemoryPool(hashId)) {
            *pPoolPtr = poolPtr;
            return KVC_OK;
        }
    }

    *pPoolPtr = nullptr;
    auto end = getCurrentTimestampUs();
    LOG_DEBUG_FMT("cost time : {} us", end - start);
    return KVC_ERR_NOT_FOUND;
}




void KVCacheServiceImpl::init() {
    auto start = getCurrentTimestampUs();
    LOG_INFO_FMT("KVCacheService init start.");

    getThreadPool.init();
    putThreadPool.init();

    auto new_pool = std::make_unique<MemoryPool>();
    LOG_INFO_FMT("createMemoryPool MEMPOOL_SIZE : {},  BLOCK_SIZE : {}",
                 MEMPOOL_SIZE, BLOCK_SIZE);
    
    // TODO(yangxianpku: 202\5.08.26): 一级指针是否有问题
    const int32_t ret = createMemoryPool(MEMPOOL_SIZE, 
                                        BLOCK_SIZE, 
                                        true, 
                                        true, 
                                        new_pool.get()
                                    );
    _pools.emplace(BLOCK_SIZE, std::move(new_pool));

    LOG_DEBUG_FMT("KVCacheService initialized. {} ms",
                  (getCurrentTimestampUs() - start) /
                      MICROSECONDS_PER_MILLISECOND);

    // 初始化storage client
    MemoryPool *pool = nullptr;
    getMemoryPool(BLOCK_SIZE, &pool);
    if (ret != KVC_OK) {
        LOG_ERROR_FMT("KVCacheServiceImpl init createMemoryPool failed!");
        return;
    }
    LOG_DEBUG_FMT("newpool:{}, pool:{}, shmpath:{}, shmAddr:{}, mountPoint:{}",
                  new_pool.get(), pool, pool->getShmFilePath().c_str(),
                  pool->getMemStartAddr(), MOUNT_POINT.c_str());
#if defined(USE_POSIX)
    storageClient.init(FileSystem::POSIX, pool->getShmFilePath(),
                       pool->getMemStartAddr(), MOUNT_POINT);
#elif defined(USE_THREEFS)
    storageClient.init(FileSystem::THREEFS, pool->getShmFilePath(),
                       pool->getMemStartAddr(), MOUNT_POINT);
#endif
    isReady = true;

#ifdef USE_DYNAMO
    _publisher = std::make_unique<DynamoPublisher>();
    _publisher->reg(pool);
    LOG_INFO_FMT("init dynamo publisher success!");
#endif

    auto end = getCurrentTimestampUs();
    LOG_INFO_FMT("KVCacheService init end.");
    LOG_DEBUG_FMT("cost time : {} us", end - start);
}




/**
 *! @brief 删除gRPC服务端，释放相关资源
 * 
 */
KVCacheServiceImpl::~KVCacheServiceImpl() {
    auto start = getCurrentTimestampUs();

    Logger::instance().shutdown();
    getThreadPool.shutdown();
    putThreadPool.shutdown();

    for (auto &pair : _pools) {
        if (pair.second) {
            destroyMemoryPool(pair.second.get());
        }
    }
    _pools.clear();

    auto end = getCurrentTimestampUs();
    LOG_DEBUG_FMT("cost time : {} us", end - start);
}